#include "multi_camera_cooperation/Assignment/linear_sum_assignment.h"

using namespace std;

void linear_sum_assignment(const vector<cv::Point2f> &marker_pixels_left, const vector<cv::Point2f> &marker_pixels_right, vector<cv::Point2f> &marker_pixels_left_buffer, vector<cv::Point2f> &marker_pixels_right_buffer, vector<vector<double>> &match_evalution, double distance_threshold){
    // initialize the marker_pixels_left_buffer and marker_pixels_right_buffer
    marker_pixels_left_buffer.clear();
    marker_pixels_right_buffer.clear();
    for(int i = 0; i < marker_pixels_left.size(); i++){
        marker_pixels_left_buffer.push_back(marker_pixels_left[i]);
    }
    for(int j = 0; j < marker_pixels_right.size(); j++){
        marker_pixels_right_buffer.push_back(marker_pixels_right[j]);
    }

    // change the size of match_evalution
    vector<int> erase_row = {};
    vector<int> erase_col = {};
    for(int i = 0; i < marker_pixels_left.size(); i++){
        double min_distance = numeric_limits<double>::max();
        for(int j = 0; j < marker_pixels_right.size(); j++){
            if(match_evalution[i][j] < min_distance){
                min_distance = match_evalution[i][j];
            }
        }
        if(min_distance > distance_threshold){
            erase_row.push_back(i);
        }
    }
    for(int j = 0; j < marker_pixels_right.size(); j++){
        double min_distance = numeric_limits<double>::max();
        for(int i = 0; i < marker_pixels_left.size(); i++){
            if(match_evalution[i][j] < min_distance){
                min_distance = match_evalution[i][j];
            }
        }
        if(min_distance > distance_threshold){
            erase_col.push_back(j);
        }
    }

    if(erase_row.size() == marker_pixels_left.size()){
        erase_col.clear();
    }

    for(int i = erase_row.size() - 1; i >= 0; i--){
        match_evalution.erase(match_evalution.begin() + erase_row[i]);
        marker_pixels_left_buffer.erase(marker_pixels_left_buffer.begin() + erase_row[i]);
    }
    for(int j = erase_col.size() - 1; j >= 0; j--){
        for(int i = 0; i < match_evalution.size(); i++){
            match_evalution[i].erase(match_evalution[i].begin() + erase_col[j]);
        }
        marker_pixels_right_buffer.erase(marker_pixels_right_buffer.begin() + erase_col[j]);
    }

    HungarianOptimizer<double> optimizer;
    std::vector<std::pair<size_t, size_t>> assignments;

    UpdateCosts(match_evalution, optimizer.costs());

    // entry of hungarian optimizer minimum-weighted matching
    optimizer.Minimize(&assignments);

    // update the marker_pixels_left_buffer and marker_pixels_right_buffer
    vector<cv::Point2f> matched_marker_pixels_left, matched_marker_pixels_right;
    for (const auto& assignment : assignments) {
        matched_marker_pixels_left.push_back(marker_pixels_left_buffer[assignment.first]);
        matched_marker_pixels_right.push_back(marker_pixels_right_buffer[assignment.second]);
    }
    marker_pixels_left_buffer = matched_marker_pixels_left;
    marker_pixels_right_buffer = matched_marker_pixels_right;

    std::cout << "The assignments result size: " << matched_marker_pixels_left.size() << std::endl;

    // release the memory of HungarianOptimizer
    optimizer.OptimizationClear();

}

void UpdateCosts(const std::vector<std::vector<double>>& association_mat, SecureMat<double>* costs) {
  size_t rows_size = association_mat.size();
  size_t cols_size = rows_size > 0 ? association_mat.at(0).size() : 0;

  costs->Resize(rows_size, cols_size);

  for (size_t row_idx = 0; row_idx < rows_size; ++row_idx) {
    for (size_t col_idx = 0; col_idx < cols_size; ++col_idx) {
      (*costs)(row_idx, col_idx) = association_mat.at(row_idx).at(col_idx);
    }
  }
}

void PrintAssignments(const std::vector<std::pair<size_t, size_t>>& assignments) {
  std::cout << "\nThe assignments result are: \n" << std::endl;

  for (const auto& assignment : assignments) {
    std::cout << "    (" << assignment.first << ", " << assignment.second << ")"
              << std::endl;
  }

  std::cout << std::endl;
}